#!/bin/bash
# Copyright 2024 Flant JSC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

mkdir -p /opt/deckhouse/bin
cp -f xtables-legacy-multi /opt/deckhouse/bin
cp -f xtables-nft-multi /opt/deckhouse/bin

if [ -f /opt/deckhouse/bin/iptables ]; then
  # If iptables symlinks already exist, we don't update them.
  exit 0
fi

NFT_KERNEL_LIMIT=3.14 # nf_tables_ipv6, nf_tables_bridge, nf_tables_arp allow since Linux kernel >= 3.14.
CURRENT_KERNEL_VERSION=$(uname -r | awk -F"-" '{print $1}') #version kernel view format: a.b.c-generic to a.b.c.

function checkNFTSupportInKernelBuildConfiguration() {
  local configNFTablesRegex='CONFIG_(NF|NFT)_(TABLES|TABLES_IPV4|TPROXY|SOCKET)=(y|m)'
  local configFiles=(
    /boot/config-$(uname -r)
    /proc/config.gz
    /lib/modules/$(uname -r)/build/.config
  )

  for file in "${configFiles[@]}"; do
    if [ -f "$file" ]; then
      local matchCount
      if [[ "$file" == *.gz ]]; then
        matchCount=$(cat "$file" | gzip -d | grep -Ec "$configNFTablesRegex" )
      else
        matchCount=$(grep -Ec "$configNFTablesRegex" "$file")
      fi
      
      if [[ $matchCount -ge 4 ]]; then
        # CONFIG_NF_TABLES, CONFIG_NF_TABLES_IPV4, 
        # CONFIG_NFT_SOCKET, CONFIG_NFT_TPROXY
        # are present for nft support
        return 0
      else
        return 1
      fi
    fi
  done

  # If no kernel configuration files were found
  return 1
}

function checkNFTSupportByAddingNFTRules() {
  local iptablesCommand="/opt/deckhouse/bin/xtables-nft-multi iptables"
  local rule="-t filter -A OUTPUT -s 127.216.216.216 -d 127.216.216.216 -j ACCEPT"

  # Check if the rule can be added
  if $($iptablesCommand $rule); then
    # Delete the rule if it was successfully added
    $iptablesCommand -t filter -D OUTPUT -s 127.216.216.216 -d 127.216.216.216 -j ACCEPT
    return 0
  fi

  return 1
}

function check_python() {
  for pybin in python3 python2 python; do
    if command -v "$pybin" >/dev/null 2>&1; then
      python_binary="$pybin"
      return 0
    fi
  done
  echo "Python not found, exiting..."
  return 1
}

function isLegacyKernel() {
  check_python
  $python_binary -c "exit(0) if tuple(map(int, '$CURRENT_KERNEL_VERSION'.split('.'))) < tuple(map(int, '$NFT_KERNEL_LIMIT'.split('.')))  else exit(1)"
}

function isNftSupport() {
  if ( checkNFTSupportInKernelBuildConfiguration ) && ( checkNFTSupportByAddingNFTRules ) ; then
    return 0
  else
    return 1
  fi
}

kubeletChainsRegex='^:(KUBE-IPTABLES-HINT|KUBE-KUBELET-CANARY)'
IPTABLES_LEGACY_RULE=$( (/opt/deckhouse/bin/xtables-legacy-multi iptables-save || true ) 2>/dev/null | grep  -E ${kubeletChainsRegex} | wc -l )

if [[ ${IPTABLES_LEGACY_RULE} -ne 0 ]] | (isLegacyKernel); then
  iptablesModeBin="xtables-legacy-multi"
elif isNftSupport; then
  iptablesModeBin="xtables-nft-multi"
else
  iptablesModeBin="xtables-legacy-multi"
fi

cmds="iptables iptables-save iptables-restore ip6tables ip6tables-save ip6tables-restore"
for cmd in $cmds; do
  ln -sf /opt/deckhouse/bin/${iptablesModeBin} /opt/deckhouse/bin/${cmd}
done
